package com.atcyj.gmall.realtime.dim.function;

import com.alibaba.fastjson.JSONObject;
import com.atcyj.gamll.realtime.common.bean.TableProcessDim;
import com.atcyj.gamll.realtime.common.constant.Constant;
import com.atcyj.gamll.realtime.common.util.JdbcUtil;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.util.HashMap;
import java.util.List;

/**
 * @author: cyj
 * @date: 2024/8/15
 */
public class DimBroadcastFunction extends BroadcastProcessFunction<JSONObject, TableProcessDim, Tuple2<JSONObject, TableProcessDim>> {
    private HashMap<String, TableProcessDim> hashMap = new HashMap<>();
    private MapStateDescriptor<String, TableProcessDim> broadcastState;

    public DimBroadcastFunction(MapStateDescriptor<String, TableProcessDim> broadcastState) {
        this.broadcastState = broadcastState;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        // 如果主流数据来了，广播数据还没来得话，从上下文中获取的广播状态中没有数据，判断会有误
        // 预加载初始的维度配置表数据
        java.sql.Connection mysqlConnection = JdbcUtil.getMysqlConnection();
        String sql = "select * from " + Constant.MYSQL_PROCESS_DATABASE + "." + Constant.MYSQL_PROCESS_TABLE + ";";
        List<TableProcessDim> tableProcessDims = JdbcUtil.queryList(mysqlConnection, sql, TableProcessDim.class, true);
        for (TableProcessDim tableProcessDim : tableProcessDims) {
            tableProcessDim.setOp("r");
            hashMap.put(tableProcessDim.getSourceTable(), tableProcessDim);
        }
        JdbcUtil.closeConnection(mysqlConnection);
    }

    /**
     * 数据流的处理方法： 数据流 只能 读取 广播状态，不能修改
     * @param jsonObject
     * @param ctx  只读的
     * @param out
     * @throws Exception
     */
    @Override
    public void processElement(JSONObject jsonObject, ReadOnlyContext ctx, Collector<Tuple2<JSONObject, TableProcessDim>> out) throws Exception {
        // 处理主流数据

        //通过上下文获取广播状态，取出里面的值（只读，不能修改）
        ReadOnlyBroadcastState<String, TableProcessDim> tableProcessState = ctx.getBroadcastState(broadcastState);

        // 查询广播状态，判断当前数据是否为维度表
        String table = jsonObject.getString("table");
        TableProcessDim tableProcessDim = tableProcessState.get(table);
        // 如果主流数据来早了，状态为空时，使用预加载的数据进行判断
        if (tableProcessDim == null) {
            tableProcessDim = hashMap.get(table);
        }
        if (tableProcessDim != null) {
            // 状态不为空，当前数据时维度表数据
            out.collect(Tuple2.of(jsonObject, tableProcessDim));
        }
    }

    /**
     * 广播后的配置流的处理方法:  只有广播流才能修改 广播状态
     *
     * @param dim
     * @param ctx
     * @param out
     * @throws Exception
     */
    @Override
    public void processBroadcastElement(TableProcessDim dim, Context ctx, Collector<Tuple2<JSONObject, TableProcessDim>> out) throws Exception {
        // 处理广播流数据,将广播流数据存到状态中
        // 通过上下文获取广播状态，往里面写数据
        BroadcastState<String, TableProcessDim> tableProcessDimState = ctx.getBroadcastState(broadcastState);
        if ("d".equals(dim.getOp())) {
            tableProcessDimState.remove(dim.getSourceTable());
            hashMap.remove(dim.getSourceTable());
        } else {
            tableProcessDimState.put(dim.getSourceTable(), dim);
        }

    }
}
